import math

# ------------------------
# 计算欧氏距离
# ------------------------
def euclidean_distance(a, b):
    return math.sqrt(sum((a[i] - b[i]) ** 2 for i in range(len(a))))

# ------------------------
# 找到距离最近的 K 个邻居
# ------------------------
def get_k_neighbors(train_data, train_labels, test_sample, k):
    distances = []

    # 计算 test_sample 到所有训练数据的距离
    for x, label in zip(train_data, train_labels):
        dist = euclidean_distance(test_sample, x)
        distances.append((dist, label))

    # 按距离升序排序
    distances.sort(key=lambda x: x[0])

    # 返回距离最近的 k 个
    return distances[:k]

# ------------------------
# 多数投票
# ------------------------
def majority_vote(labels):
    count_dict = {}
    for lab in labels:
        count_dict[lab] = count_dict.get(lab, 0) + 1

    # 返回出现次数最多的标签
    return max(count_dict, key=count_dict.get)

# ------------------------
# KNN 预测单个样本
# ------------------------
def KNN_predict(train_data, train_labels, test_sample, k):
    neighbors = get_k_neighbors(train_data, train_labels, test_sample, k)
    labels = [label for _, label in neighbors]
    return majority_vote(labels)

# ------------------------
# KNN 预测多个样本
# ------------------------
def KNN_classify(train_data, train_labels, test_data, k):
    predictions = []
    for sample in test_data:
        predictions.append(KNN_predict(train_data, train_labels, sample, k))
    return predictions

# ===========================================
# 示例数据
# ===========================================
train_X = [
    [1, 2], [2, 3], [3, 3],
    [8, 8], [9, 8], [8, 9]
]

train_y = [0, 0, 0, 1, 1, 1]

test_X = [
    [2, 2],
    [8, 7],
    [10, 10]
]

result = KNN_classify(train_X, train_y, test_X, k=3)
print(result)
